# -*- coding: UTF-8 -*-
# @Project -> File   ：AIserving -> pipeline
# @Author ：luoyang
# @Date   ：2021/3/2 10:52

import os
import sys

import numpy as np

from confs.engine_configs import engs
from nlp_tools.tasks.abs_task_model import ABCTaskModel
from project_settings import NLP_TOOLS_MODELS_PATH
from nlp_tools.utils.ner_utils import format_ner_result

#from processing import SentenceEncoder, AfterProcess


_global_engine_dict = {}

def get_engine(eng_name, model_name,use_grpc):

    """
    获取推理引擎实例
    :param eng_name: string，引擎名称，目前支持：tfserving，onnxserving
    :param model_name: string，模型名称

    使用：
        model = get_engine('tfserving', 'direction')
        res = model.predict(inputdata)

    """
    eng_class = engs[eng_name]['engine']
    eng_info = engs[eng_name]['urls'][model_name]
    eng_obj = eng_class(tfserving_dict=eng_info,use_grpc=use_grpc)

    return eng_obj

class PipeLine(object):

    def __init__(self,model_name,engine_type='tfserving',use_grpc=True):
        self.ai_engine = get_engine(engine_type, model_name,use_grpc)
        self.model_name = model_name

        self.nlp_model = ABCTaskModel.load_model(model_path=self.get_nlp_tools_model_path(),no_models=True)
        #self.after_pro = AfterProcess()

    def get_nlp_tools_model_path(self):
        return os.path.join(NLP_TOOLS_MODELS_PATH,self.model_name)

    def predict(self, sentences):

        """
        实体预测整体流程，包含字符串转token，模型预测，处理模型预测结果等
        :param sentences: list，每个元素为一个字符串，表示待预测的一句话，形如：
                            ['我爱中国', '南京是江苏的省会', ……]
        :return:
            is_success: bool，表示是否成功用模型进行预测，如果此项为True，res为空，则说明成功预测，但没有预测出包含实体
            res: list，元素为list，包含每条输入字符串的预测结果，顺序和输入sentences保持一致，形如：
                 [[[0, 2, '南京', 'LOC'], [5, 8, '中华中学', 'ORG'], ……],
                  [[0, 3, '刘德华', 'PER'], [10, 12, '香港', LOC], ……],
                  ……]
        """

        tensor,segment_texts = self.nlp_model.get_tfserver_inputs(sentences)
        engine_predicts = self.ai_engine.predict(tensor)

        predict_labels = self.nlp_model.get_tfserver_http_result_labels(sentences,engine_predicts)
        res = format_ner_result(predict_labels,segment_texts,sentences)
        is_success = True

        return is_success, res
